import torch
import torch.nn as nn
import numpy as np


class ZhnNetHead(nn.Module):
    def __init__(self, classify=False, device=torch.device('cpu')):
        super().__init__()
        self.conv1 = nn.Conv2d(96, 5, kernel_size=3, padding=1, bias=True)
        self.conv2 = nn.Conv2d(48, 5, kernel_size=3, padding=1, bias=True)
        self.conv3 = nn.Conv2d(24, 5, kernel_size=3, padding=1, bias=True)
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Linear(192, 1, bias=False)
        self.classify = classify
        self.device = device

    def forward(self, x):
        if self.classify:  # Nx192x15x20
            x = x[2]
            x = self.maxpool(x)  # Nx192x1x1
            x = x.flatten(1)  # Nx192
            x = self.fc(x)  # Nx1
            x = torch.sigmoid(x.squeeze())  # N
        else:
            x1, x2, x3 = x
            x1 = self.conv1(x1).view(-1, 5, 300)
            x2 = self.conv2(x2).view(-1, 5, 1200)
            x3 = self.conv3(x3).view(-1, 5, 4800)
            x = torch.cat([x1, x2, x3], dim=2)
        return x
